# 
from exp import data, run
import random, os
from datetime import datetime

def _examples_repr(examples):
	def _make_example(e):
		problem = e['problem']
		solution = e['cot']
		answer = e['solution']
		return f'**Problem**: {problem}\n\n**Solution**:\n\n{solution}\nThe answer is: {answer}'
	examples_ = [_make_example(e) for e in examples]
	return '\n\n'.join([f'## Example {i+1}\n\n{e}' for i,e in enumerate(examples_)])

class Kshot_Reasoner:
	def __init__(self, examples, k=0):
		self.examples = examples
		self.k = k

	def reason(self, model_id, question, verbal=True):
		random.shuffle(self.examples)
		examples = [e for e in self.examples[:self.k+1] if e['problem'] != question]
		examples = _examples_repr(examples[:self.k])
		resp = run.generate(model_id, question, 'com/reasoner_prompts', 'k_shot',
							question=question.question, 
							examples=examples, 
							verbal=verbal,
							)
		return question.extra_asnwer(resp), 'The answer is', []

def test(test_set):
	math500 = data.get_data('MATH500')
	test_examples = data.get_data(test_set)
	reasoner = Kshot_Reasoner(math500, k=8)
	for e in test_examples[:10]:
		run.test(e, reasoner.reason)

out_path = '../OUT/k_shot'

def exp(test_set):
	math500 = data.get_data('MATH500')
	test_examples = data.get_data(test_set)
	for k in [8]:
		reasoner = Kshot_Reasoner(math500, k=k)
		outfile = f'{out_path}/{test_set}_{k}-shot.jsonl'
		print(f'# {k} {outfile} {datetime.now()}')
		if os.path.isfile(outfile): continue
		run.run_exp(test_examples, reasoner.reason,  num_samples=8, max_retries=8, outfile=outfile)

def main():
	if not os.path.isdir(out_path):
		os.mkdir(out_path)
	exp('MATH500')
	exp('AIME')

if __name__ == '__main__':
	main()
